package com.ruoyi.ai.factory.impl;

import com.alibaba.dashscope.common.Role;
import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.core.builder.ChatBuilder;
import com.google.common.collect.Lists;
import com.ruoyi.ai.domain.vo.WxChatVo;
import com.ruoyi.ai.domain.vo.WxModelConfigVo;
import com.ruoyi.ai.enums.MfgEnum;
import com.ruoyi.ai.factory.ModelChatTemplate;
import com.ruoyi.ai.ws.ChatWebSocketManager;
import com.ruoyi.common.exception.ServiceException;
import io.github.lnyocly.ai4j.listener.SseListener;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletion;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatMessage;
import io.github.lnyocly.ai4j.service.IChatService;
import io.github.lnyocly.ai4j.service.PlatformType;
import io.github.lnyocly.ai4j.service.factor.AiService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 百度千帆
 */
@Slf4j
@Service
public class OllamaModelChatService extends ModelChatTemplate {
    @Autowired
    private AiService aiService;

    public OllamaModelChatService() {
        super(MfgEnum.OLLAMA.getType());
    }

    @Resource
    private ChatWebSocketManager chatWebSocketManager;

    @Override
    public String chatStream(List<WxChatVo> chatVoList, Long userId) {
        // 获取chat服务实例
        IChatService chatService = aiService.getChatService(PlatformType.OLLAMA);

        // 构造基本请求参数
        ChatCompletion chatCompletion = getChatCompletionBase();
        // 构造消息参数
        chatCompletion.setMessages(buildChatMessageList(chatVoList));
        // 构造监听器
        SseListener sseListener = new SseListener() {
            @Override
            protected void send() {
                chatWebSocketManager.sendOneMessage(userId, this.getCurrStr());
                System.out.println(this.getCurrStr());
            }
        };
        // 显示函数参数，默认不显示
        sseListener.setShowToolArgs(true);

        // 发送SSE请求
        try {
            chatService.chatCompletionStream(chatCompletion, sseListener);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        System.out.println(sseListener.getOutput());
        //发送对话结束标记，前端好判断本次对话已结束，允许下一轮对话
        chatWebSocketManager.sendOneMessage(userId, OVER_FLAG);
        return sseListener.getOutput().toString();
    }

    @Override
    public String chatStreamOnce(String content) {
        List<WxChatVo> chatVoList = Lists.newArrayList();
        WxChatVo wxChatVo = new WxChatVo();
        wxChatVo.setRole(Role.USER.getValue());
        wxChatVo.setContent(content);
        chatVoList.add(wxChatVo);
        return chatStream(chatVoList, -1L);
    }

    private ChatCompletion getChatCompletionBase() {
        WxModelConfigVo config = getConfig();
        return ChatCompletion.builder()
                .model(config.getModelType())
                .build();
    }

    private List<ChatMessage> buildChatMessageList(List<WxChatVo> chatVoList) {
        return chatVoList.stream().map(chatVo -> {
            if (chatVo.getRole().equals(Role.USER.getValue())) {
                return ChatMessage.withUser(chatVo.getContent());
            } else if (chatVo.getRole().equals(Role.ASSISTANT.getValue())) {
                return ChatMessage.withAssistant(chatVo.getContent());
            } else {
               throw new ServiceException("不支持的role类型");
            }
        }).collect(Collectors.toList());
    }

}
